from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate
from langchain_core.example_selectors import SemanticSimilarityExampleSelector
from langchain_community.vectorstores import FAISS
from patagent.llm import PatsnapLLM
from patagent.prompt import PARSE_OCR_RESULT_FEWSHOT_PROMPT
from langchain_openai import OpenAIEmbeddings
from patagent.constant import (
    PATSNAP_API_URL,
    PATSNAP_API_KEY,
    OPENAI_MODEL_HEADERS
)


llm = PatsnapLLM()

examples = [
    {"input": '请从输入的内容中提取`公司名称`，输入：\n\n$$\npm textcircled {<} \times textcircled {<} \times times textcircled {<} \textcircled{<} \downarrow\downarrow\downarrow\cup P 1 0 0\n$$\n上海联影医疗科技股份有限公司\n展讯通信(上海)有限公司\n上海商汤智能科技有限公司\n中国航发商用航空发动机有限责任公司', "output": '["上海联影医疗科技股份有限公司","展讯通信(上海)有限公司","上海商汤智能科技有限公司","中国航发商用航空发动机有限责任公司"]'},
    {"input": '请从输入的内容中提取`公司名称`，输入：\n\n企业名称\n总部所在省份\n参评赛道\n9爱德德曼氢能\n浙江省\n氢能\n奥扬科技\n山东省\n氢能\n田锋源氢能', "output": '["爱德德曼氢能","奥扬科技","田锋源氢能"]'},
    {"input": '请从输入的内容中提取`公司名称`，输入：\n\n## 福布斯 Forbes\n2024 福布斯中国 $\boldsymbol{\Lambda}$ 、二智能科技企业\nTODRA\n1阿里云阿里云计算有限公司\n ${\bf2}$ 爱诗科技北京爱诗科技有限公 $\mathbf{3}$ 爱芯元智爱芯元智半导体股份有限公司 $\boldsymbol{4}$ 百川智能北京百川智能科技有限公', "output": '["阿里云计算有限公司","北京爱诗科技有限公司","爱芯元智半导体股份有限公司","北京百川智能科技有限公司"]'},
    {"input": '请从输入的内容中提取`排名`，输入：\n\n排名|  公司名称 |  行业\n---|---|---\n_1_ |  北京瓦特曼智能科技有限公司 | 智能制造及服务\n_2_ |  图湃（北京）医疗科技有限公司 | 药物和医疗器械研发及商业化服务\n_3_ |  北京中科原动力科技有限公司 | 农业现代化\n_4_ |  北京云圣智能科技有限责任公司 | 智能制造及服务\n_5_ |  杭州纤纳光电科技股份有限公司 | 绿色科技', "output": '["1","2","3","4","5"]'},
    {"input": '请从输入的内容中提取`品牌名称或简称`，输入：\n\n| 排名 | 公司名称                                   | 品牌名(或简称)            | 近三年收入增长率 |\n| ---- | ------------------------------------------ | ------------------------- | ---------- |\n| 1    | 广州市莱檬生物科技有限公司                       | 莱檬生物、优莱檬                    | 3534%      |\n| 2    | 广州蓝海机器人系统有限公司                             | LH-AGV、LH-BZX                       | 1790%      |\n| 3    | 深圳鲲云信息科技有限公司                             | Corerain, CAISA                    | 1631%      |', "output": '["莱檬生物、优莱檬","LH-AGV、LH-BZX","Corerain, CAISA"]'},
    {"input": '请从输入的内容中提取`增长率`，输入：\n\n| 排名 | 公司名称                                   | 品牌名(或简称)            | 近三年收入增长率 |\n| ---- | ------------------------------------------ | ------------------------- | ---------- |\n| 1    | 中昊芯英(杭州)科技有限公司                       | 中昊芯英                    | 3534%      |\n| 2    | 镁佳北京科技有限公司                             | 镁佳                       | 1790%      |\n| 3    | 深圳作为科技有限公司                             | ZUOWEI                    | 1631%      |', "output": '["3534%","1790%","1631%"]'},
    {"input": '请从输入的内容中提取`城市或省份`，输入：\n\n公司名称 | 品牌名(或简称) | 城市\n--- | --- | ---\n广州百吉生物制药有限公司 | BioSyngen | 广州\n凯米生物医药(成都)有限公司 | ChimigenBio、SynNeogen | 成都\n航天时代飞鹏有限公司 | Flyingpon、飞吉送 | 昆山', "output": '["广州","成都","昆山"]'},
    {"input": '请从输入的内容中提取`城市或省份`，输入：\n\n企业名称\n总部所在省份\n参评赛道\n爱德德曼氢能\n浙江省\n氢能\n奥扬科技\n山东省\n氢能\n锋源氢能\n北京市\n氢能\n国氢科技\n北京市\n氢能', "output": '["浙江省","山东省","北京市","北京市"]'},
    {"input": '请从输入的内容中提取`行业或赛道或领域`，输入：\n\n排名|  公司名称 |  赛道\n---|---|---\n_1_ |  华为投资控股有限公司 | 电信、通讯和计算机软硬件\n_2_ |  比亚迪股份有限公司 | 高端制造\n_3_ |  宁德时代新能源科技股份有限公司 | 新能源及新材料', "output": '["电信、通讯和计算机软硬件","高端制造","新能源及新材料"]'},
    {"input": '请从输入的内容中提取`行业或赛道或领域`，输入：\n\n排名|  公司名称 |  行业  \n---|---|---  \n_1_ |  北京瓦特曼智能科技有限公司 | 智能制造及服务  \n_2_ |  图湃（北京）医疗科技有限公司 | 药物和医疗器械研发及商业化服务  \n_3_ |  北京中科原动力科技有限公司 | 农业现代化  ', "output": '["智能制造及服务","药物和医疗器械研发及商业化服务","农业现代化"]'},
    {"input": '请从输入的内容中提取`利润`，输入：\n\n排名 | 上年 _排名_ | 公司名称 _(中文)_ | 营业收入 _(百万美元)_ | 利润 _(百万美元)_ | _关键_ _数据_  \n---|---|---|---|---|---  \n1| 1| 国家电网有限公司 | 545,947.5| 9,204.3| +  \n2| 3| 中国石油化工集团有限公司 | 429,699.7| 9,393.4| +  \n3| 2| 中国石油天然气集团有限公司 | 421,713.6| 21,294.7| +  ', "output": '["9204.3","9393.4","21294.7"]'},
    {"input": '请从输入的内容中提取`营业收入`，输入：\n\n排名 | 上年 _排名_ | 公司名称 _(中文)_ | 营业收入 _(百万美元)_ | 利润 _(百万美元)_ | _关键_ _数据_  \n---|---|---|---|---|---  \n1| 1| 国家电网有限公司 | 545,947.5| 9,204.3| +  \n2| 3| 中国石油化工集团有限公司 | 429,699.7| 9,393.4| +  \n3| 2| 中国石油天然气集团有限公司 | 421,713.6| 21,294.7| +  ', "output": '["545947.5","429699.7","421713.6"]'},
    {"input": '请从输入的内容中提取`企业价值`，输入：\n\n| 排名 | 企业 | 企业价值 (亿元人民币） | 细分领域 | 具体优势 | 总部城市 |\n| --- | --- | --- | --- | --- | --- |\n| 1 | 寒武纪 | 2380 | 算力硬件类 | AI芯片 | 北京 |\n| 2 | 科大讯飞 | 1160 | 语音识别类 | 语音识别/自然语言处理 | 安徽合肥 |\n| 3 | 商汤科技 | 500 | 视觉识别类 | 视觉识别及AIGC | 上海 |\n| 4 | 地平线 | 430 | 算力硬件类 | 自动驾驶芯片及算法 | 北京 |\n| 5 | 小马智行 | 320 | 自动驾驶类 | 自动驾驶算法 | 广东广州 |', "output": '["2380","1160","500","430","320"]'}
]
to_vectorize = [" ".join(example.values()) for example in examples]
embeddings = OpenAIEmbeddings(openai_api_base=PATSNAP_API_URL, openai_api_key=PATSNAP_API_KEY, default_headers=OPENAI_MODEL_HEADERS)
vectorstore = FAISS.from_texts(to_vectorize, embeddings, metadatas=examples)

example_selector = SemanticSimilarityExampleSelector(
    vectorstore=vectorstore,
    k=1,
)

# Prompt
example_prompt = ChatPromptTemplate.from_messages(
    [
        ("human", "{input}"),
        ("ai", "{output}")
    ]
)
few_shot_prompt = FewShotChatMessagePromptTemplate(
    # The input variables select the values to pass to the example_selector
    input_variables=["question"],
    example_selector=example_selector,
    # Define how each example will be formatted.
    # In this case, each example will become 2 messages:
    # 1 human, and 1 AI
    example_prompt=example_prompt,
)

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", PARSE_OCR_RESULT_FEWSHOT_PROMPT),
        few_shot_prompt,
        ("human", "{question}"),
    ]
)

# Chain
fewshot_extractor = prompt | llm | StrOutputParser()